import numpy as np
import pandas as pd
from scipy.fft import fft, fftfreq
from scipy import signal
import matplotlib.pyplot as plt

# 设置全局字体为支持中文的字体，如黑体 'SimHei'
plt.rcParams['font.sans-serif'] = ['SimHei']  # 或者 ['Microsoft YaHei']
# 解决保存图像时负号 '-' 显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False


def extract_frequency_domain_features(signal_data, sampling_rate):
    """
    提取信号的频域特征

    参数:
    signal_data: 一维数组，输入的时间序列信号数据
    sampling_rate: 采样率（Hz）

    返回:
    features_dict: 字典，包含提取的各个频域特征值
    freq_spectrum: 频率谱数据，可用于可视化
    """
    # 确保输入为numpy数组
    signal_data = np.array(signal_data)
    n = len(signal_data)

    # 计算FFT（快速傅里叶变换）
    fft_values = fft(signal_data)
    frequencies = fftfreq(n, d=1 / sampling_rate)

    # 取单侧频谱（正频率部分）
    positive_freq_idx = frequencies > 0
    positive_frequencies = frequencies[positive_freq_idx]
    positive_fft = np.abs(fft_values[positive_freq_idx]) / n  # 归一化振幅
    positive_fft[1:] *= 2  # 除直流分量外乘以2

    # 计算功率谱密度 (PSD)
    frequencies_psd, psd = signal.welch(signal_data, fs=sampling_rate, nperseg=min(256, n // 4))

    # 预先计算频谱重心，避免在字典未完全构建前引用它
    spectral_centroid = np.sum(positive_frequencies * positive_fft) / np.sum(positive_fft)

    # 计算频谱带宽
    spectral_bandwidth = np.sqrt(
        np.sum((positive_frequencies - spectral_centroid) ** 2 * positive_fft) / np.sum(positive_fft))

    # 提取频域特征
    features_dict = {
        'dominant_frequency': positive_frequencies[np.argmax(positive_fft)],  # 主导频率
        'max_amplitude': np.max(positive_fft),  # 最大振幅
        'mean_amplitude': np.mean(positive_fft),  # 平均振幅
        'amplitude_std': np.std(positive_fft),  # 振幅标准差
        'spectral_energy': np.sum(positive_fft ** 2),  # 频谱能量
        'spectral_centroid': spectral_centroid,  # 频谱重心
        'spectral_bandwidth': spectral_bandwidth,  # 频谱带宽
        'spectral_entropy': -np.sum(positive_fft * np.log2(positive_fft + 1e-10)),  # 频谱熵（加小值避免log(0)）
        'psd_max': np.max(psd),  # 最大功率谱密度
        'psd_mean': np.mean(psd),  # 平均功率谱密度
        'frequency_ratio': positive_frequencies[np.argmax(positive_fft)] / (sampling_rate / 2)  # 频率比（主导频率与奈奎斯特频率之比）
    }

    # 准备频谱数据用于可视化
    freq_spectrum = {
        'frequencies': positive_frequencies,
        'amplitudes': positive_fft,
        'frequencies_psd': frequencies_psd,
        'psd': psd
    }

    return features_dict, freq_spectrum


# 示例用法
if __name__ == "__main__":
    # 生成示例信号（含噪声的正弦波）
    sampling_rate = 1000  # 采样率1000Hz
    t = np.linspace(0, 1, sampling_rate)  # 1秒时间序列
    signal_data = 0.8 * np.sin(2 * np.pi * 50 * t) + 0.2 * np.sin(2 * np.pi * 120 * t) + 0.1 * np.random.normal(
        size=len(t))

    # 提取频域特征
    features, spectrum = extract_frequency_domain_features(signal_data, sampling_rate)

    # 打印特征值
    print("频域特征提取结果:")
    for key, value in features.items():
        print(f"{key}: {value:.6f}")

    # 可视化信号和频谱
    plt.figure(figsize=(15, 10))

    # 绘制原始信号
    plt.subplot(2, 2, 1)
    plt.plot(t, signal_data)
    plt.title('时域信号')
    plt.xlabel('时间 (s)')
    plt.ylabel('振幅')
    plt.grid(True)

    # 绘制振幅频谱
    plt.subplot(2, 2, 2)
    plt.plot(spectrum['frequencies'], spectrum['amplitudes'])
    plt.title('振幅频谱')
    plt.xlabel('频率 (Hz)')
    plt.ylabel('振幅')
    plt.grid(True)

    # 标记主导频率
    dominant_freq = features['dominant_frequency']
    max_amp = features['max_amplitude']
    plt.plot(dominant_freq, max_amp, 'ro', label=f'主导频率: {dominant_freq:.1f}Hz')
    plt.legend()

    # 绘制功率谱密度 (PSD)
    plt.subplot(2, 2, 3)
    plt.semilogy(spectrum['frequencies_psd'], spectrum['psd'])
    plt.title('功率谱密度 (PSD)')
    plt.xlabel('频率 (Hz)')
    plt.ylabel('PSD (dB/Hz)')
    plt.grid(True)

    # 绘制特征值条形图
    plt.subplot(2, 2, 4)
    feature_names = list(features.keys())
    feature_values = list(features.values())
    plt.barh(range(len(feature_names)), feature_values)
    plt.yticks(range(len(feature_names)), feature_names)
    plt.title('频域特征值')
    plt.xlabel('特征值')
    plt.tight_layout()

    plt.show()